import numpy as np
import matplotlib.pyplot as plt
import csv
import os
import argparse

def calculate_learning_curve(y1, learning_rate, max_training=30):
    """
    计算改良学习曲线的预测分数
    :param y1: 第1次训练的分数
    :param learning_rate: 学习率 (LR)
    :param max_training: 最大训练次数 (默认30)
    :return: x轴数据, 预测分数列表
    """
    # 计算学习曲线参数b
    b = np.log(learning_rate) / np.log(2)
    
    x_values = np.arange(1, max_training + 1)
    predicted_scores = []
    
    for x in x_values:
        if x == 1:
            # 第1次直接使用y1
            zx = y1
        else:
            # 计算 Z_x = y1 * [x^(b+1) - (x-1)^(b+1)]
            term1 = np.power(x, b + 1)
            term2 = np.power(x - 1, b + 1)
            zx = y1 * (term1 - term2)
        
        predicted_scores.append(zx)
    
    return x_values, predicted_scores

def calculate_learning_curve_B(y1, learning_rate, max_training=30):
    """
    计算改良学习曲线的预测分数
    :param y1: 第1次训练的分数
    :param learning_rate: 学习率 (LR)
    :param max_training: 最大训练次数 (默认30)
    :return: x轴数据, 预测分数列表
    """
    # 计算学习曲线参数b
    b = np.log(learning_rate) / np.log(2)
    
    x_values = np.arange(1, max_training + 1)
    predicted_scores = []
    
    for x in x_values:
        if x == 1:
            # 第1次直接使用y1
            zx = y1
        else:
            # 计算 Z_x = y1 * [x^(b+1) - (x-1)^(b+1)]
            term1 = np.power(x, b)
            #term2 = np.power(x - 1, b + 1)
            zx = np.power(y1, term1)
        
        predicted_scores.append(zx)
    
    return x_values, predicted_scores


def calculating_learning_rate(y_pre):
    """
    根据预置数据计算学习率
    :param y_pre: 预先准备好备牙分数数据
    :return: 该组数据的预期学习率
    """
    x = len(y_pre)
    y_mean = np.mean(y_pre)
    LR = np.power(2, np.log10(y_mean/y_pre[0])/np.log10(x))
    return LR

def calculate_learning_rate(y_pre):
    """
    根据Wright学习曲线模型计算学习率LR
    :param y_pre: 累计平均时间/成本的列表，如 [100, 80, 70.21, 64]（对应x=1,2,3,4）
    :return: 学习率LR
    """
    # 转换为numpy数组并确保数据为正
    y_pre = np.array(y_pre)
    if np.any(y_pre <= 0):
        raise ValueError("y_pre must contain positive values")
    
    # 创建x序列 [1, 2, 3, ..., n]
    n = len(y_pre)
    x = np.arange(1, n + 1)
    
    # 取对数
    log_x = np.log10(x)
    log_y = np.log10(y_pre)
    
    # 线性回归（1次多项式拟合），返回斜率b和截距
    b, log_y1 = np.polyfit(log_x, log_y, 1)
    
    # 计算学习率 LR = 2^b
    LR = 2 ** b
    return LR

# 从CSV文件读取数据
def read_data_from_csv(file_path):
    data_groups = []
    with open(file_path, 'r') as file:
        reader = csv.reader(file)
        for row in reader:
            # 跳过空行
            if not row:
                continue
            # 将字符串转换为浮点数
            try:
                group = list(map(float, row))
                # 每组数据至少需要1个点
                if len(group) >= 1:
                    data_groups.append(group)
            except ValueError:
                continue
    return data_groups

def plot_learning_curves(data_groups, csv_file_path, max_training=30, 
                        title="多组学习曲线预测", figure_size=(12, 8), 
                        save_csv=True, save_plot=True):
    """
    绘制多组学习曲线并保存结果
    
    :param data_groups: 多组数据的列表，每组数据为一个数值列表
    :param csv_file_path: 原始CSV文件路径，用于生成输出文件名
    :param max_training: 最大训练次数，默认30
    :param title: 图表标题
    :param figure_size: 图表尺寸，默认(12, 8)
    :param save_csv: 是否保存CSV结果，默认True
    :param save_plot: 是否保存图像，默认True
    :return: 包含预测结果和有效组信息的字典
    """
    if not data_groups:
        print("未找到有效数据")
        return None
    
    # 创建图表
    plt.figure(figsize=figure_size)
    
    # 为每组数据分配不同颜色
    colors = plt.cm.tab10(np.linspace(0, 1, len(data_groups)))
    all_predictions = []
    valid_groups = []  # 存储有效组的索引和参数
    
    # 处理每组数据
    for i, group_data in enumerate(data_groups):
        y1 = group_data[0]
        try:
            LR = calculate_learning_rate(group_data)
        except ZeroDivisionError:
            print(f"组 {i+1} 数据不足，无法计算学习率 (需要至少2个点)")
            continue
        
        # 计算预测值
        x, Z = calculate_learning_curve(y1=y1, learning_rate=LR, max_training=max_training)
        all_predictions.append(Z)
        valid_groups.append((i+1, y1, LR))
        
        # 绘制预测曲线
        plt.plot(x, Z, color=colors[i], linewidth=1.5, 
                 label=f'Sample {i+1} (y1={y1:.2f}, LR={LR:.4f})')
        
        # 绘制原始数据点（可选）
        # original_x = np.arange(1, len(group_data) + 1)
        # plt.scatter(original_x, group_data[:len(original_x)], 
        #             color=colors[i], s=50, zorder=5)
    
    # 添加图例和标签
    #plt.title(title, fontsize=14, fontproperties="SimHei")
    plt.xlabel("Trial number (x)", fontsize=12, fontproperties="SimHei")
    plt.ylabel("Score", fontsize=12, fontproperties="SimHei")
    plt.grid(True, linestyle="--", alpha=0.7)
    plt.xticks(np.arange(0, max_training + 1, 5))
    plt.legend(loc='best')
    plt.ylim(bottom=None, top=100)
    
    # 保存图像
    if save_plot:
        plot_filename = os.path.splitext(csv_file_path)[0] + '_learning_curves.png'
        plt.tight_layout()
        plt.savefig(plot_filename)
        print(f"图像已保存至: {plot_filename}")
    
    plt.show()
    
    # 保存CSV结果
    if save_csv and all_predictions:
        # 创建列标题：训练次数 + 组信息
        headers = ["训练次数"]
        for group_info in valid_groups:
            group_num, y1, LR = group_info
            headers.append(f"组 {group_num} (y1={y1:.2f}, LR={LR:.4f})")
        
        # 转置数据：行->训练次数，列->不同组
        transposed_data = list(zip(range(1, max_training + 1), *all_predictions))
        
        # 写入CSV文件
        output_csv = os.path.splitext(csv_file_path)[0] + '_预测结果.csv'
        with open(output_csv, 'w', newline='', encoding='utf-8-sig') as f:
            writer = csv.writer(f)
            writer.writerow(headers)  # 写入标题行
            writer.writerows(transposed_data)  # 写入数据行
        
        print(f"预测结果已保存至: {output_csv}")
    elif not all_predictions:
        print("没有有效数据可用于保存")
    
    # 返回结果
    return {
        'predictions': all_predictions,
        'valid_groups': valid_groups,
        'x_values': list(range(1, max_training + 1))
    }

def plot_average_learning_curves(ir_csv_path, mr_csv_path, max_training=30, 
                                title="IR组与MR组平均学习曲线对比", figure_size=(12, 8), 
                                save_csv=True, save_plot=True):
    """
    绘制IR组和MR组的平均学习曲线对比图
    
    :param ir_csv_path: IR组CSV文件路径
    :param mr_csv_path: MR组CSV文件路径
    :param max_training: 最大训练次数，默认30
    :param title: 图表标题
    :param figure_size: 图表尺寸，默认(12, 8)
    :param save_csv: 是否保存CSV结果，默认True
    :param save_plot: 是否保存图像，默认True
    :return: 包含两组平均预测结果的字典
    """
    # 设置matplotlib中文字体
    plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.rcParams['axes.unicode_minus'] = False
    
    # 读取两组数据
    ir_data_groups = read_data_from_csv(ir_csv_path) if os.path.exists(ir_csv_path) else []
    mr_data_groups = read_data_from_csv(mr_csv_path) if os.path.exists(mr_csv_path) else []
    
    if not ir_data_groups and not mr_data_groups:
        print("未找到有效数据文件")
        return None
    
    # 创建图表
    plt.figure(figsize=figure_size)
    
    results = {}
    colors = ['#1f77b4', '#ff7f0e']  # 蓝色和橙色
    labels = ['IR', 'MR']
    data_groups_list = [ir_data_groups, mr_data_groups]
    file_paths = [ir_csv_path, mr_csv_path]
    
    for group_idx, (data_groups, file_path, color, label) in enumerate(zip(data_groups_list, file_paths, colors, labels)):
        if not data_groups:
            print(f"{label}数据为空，跳过")
            continue
            
        # 计算每组的预测值
        all_group_predictions = []
        valid_count = 0
        
        for i, group_data in enumerate(data_groups):
            y1 = group_data[0]
            try:
                LR = calculate_learning_rate(group_data)
                x, Z = calculate_learning_curve(y1=y1, learning_rate=LR, max_training=max_training)
                all_group_predictions.append(Z)
                valid_count += 1
            except (ZeroDivisionError, ValueError):
                print(f"{label}第{i+1}组数据无效，跳过")
                continue
        
        if not all_group_predictions:
            print(f"{label}没有有效数据")
            continue
            
        # 计算平均预测值
        # 转换为numpy数组进行平均计算
        predictions_array = np.array(all_group_predictions)
        average_predictions = np.mean(predictions_array, axis=0)
        
        # 绘制平均曲线
        x_values = np.arange(1, max_training + 1)
        plt.plot(x_values, average_predictions, color=color, linewidth=2.5, 
                 label=f'Learning curve of {label} (n={valid_count})', marker='o', markersize=4)
        
        # 可选：添加标准差阴影区域
        std_predictions = np.std(predictions_array, axis=0)
        plt.fill_between(x_values, 
                        average_predictions - std_predictions, 
                        average_predictions + std_predictions, 
                        color=color, alpha=0.2, 
                        label=f'Standard deviation range of {label}')
        
        # 保存结果
        results[label.lower()] = {
            'average_predictions': average_predictions.tolist(),
            'std_predictions': std_predictions.tolist(),
            'valid_count': valid_count,
            'x_values': x_values.tolist()
        }
    
    # 添加图例和标签
    #plt.title(title, fontsize=14, fontproperties="SimHei")
    plt.xlabel("Trial number (x)", fontsize=12, fontproperties="SimHei")
    plt.ylabel("Score", fontsize=12, fontproperties="SimHei")
    plt.grid(True, linestyle="--", alpha=0.7)
    plt.xticks(np.arange(0, max_training + 1, 5))
    
    # 设置图例中文字体
    legend = plt.legend(loc='best', fontsize=10)
    for text in legend.get_texts():
        text.set_fontproperties("SimHei")
    
    plt.ylim(bottom=None, top=100)
    
    # 保存图像
    if save_plot:
        plot_filename = 'IR_MR_average_learning_curves.png'
        plt.tight_layout()
        plt.savefig(plot_filename)
        print(f"对比图像已保存至: {plot_filename}")
    
    plt.show()
    
    # 保存CSV结果
    if save_csv and results:
        headers = ["训练次数"]
        data_rows = [list(range(1, max_training + 1))]
        
        for group_name in ['ir组', 'mr组']:
            if group_name in results:
                headers.extend([f"{group_name.upper()}平均值", f"{group_name.upper()}标准差"])
                data_rows.append(results[group_name]['average_predictions'])
                data_rows.append(results[group_name]['std_predictions'])
        
        # 转置数据
        transposed_data = list(zip(*data_rows))
        
        # 写入CSV文件
        output_csv = 'IR_MR_average_comparison.csv'
        with open(output_csv, 'w', newline='', encoding='utf-8-sig') as f:
            writer = csv.writer(f)
            writer.writerow(headers)
            writer.writerows(transposed_data)
        
        print(f"对比结果已保存至: {output_csv}")
    
    return results

# 主程序部分
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='学习曲线分析工具')
    parser.add_argument('--mode', choices=['single', 'compare'], default='single',
                        help='运行模式: single=单类预测曲线, compare=两类平均对比曲线')
    parser.add_argument('--ir_file', type=str, default='IR组.csv',
                        help='IR组数据文件路径 (默认: IR组.csv)')
    parser.add_argument('--mr_file', type=str, default='MR组.csv',
                        help='MR组数据文件路径 (默认: MR组.csv)')
    parser.add_argument('--single_file', type=str,
                        help='单类模式下的数据文件路径 (如果不指定，则使用--ir_file)')
    parser.add_argument('--max_training', type=int, default=30,
                        help='最大训练次数 (默认: 30)')
    parser.add_argument('--title', type=str,
                        help='图表标题 (如果不指定，将使用默认标题)')
    parser.add_argument('--no_save_csv', action='store_true',
                        help='不保存CSV结果文件')
    parser.add_argument('--no_save_plot', action='store_true',
                        help='不保存图像文件')
    
    args = parser.parse_args()
    
    # 根据模式执行相应功能
    if args.mode == 'single':
        # 单类预测曲线模式
        csv_file_path = args.single_file if args.single_file else args.ir_file
        
        if not os.path.exists(csv_file_path):
            print(f"错误: 文件 '{csv_file_path}' 不存在")
            exit(1)
        
        # 读取数据
        data_groups = read_data_from_csv(csv_file_path)
        
        if not data_groups:
            print("未找到有效数据")
            exit(1)
        
        # 设置标题
        if args.title:
            title = args.title
        else:
            file_name = os.path.splitext(os.path.basename(csv_file_path))[0]
            title = f"{file_name}学习曲线预测"
        
        # 调用绘图函数
        result = plot_learning_curves(
            data_groups=data_groups, 
            csv_file_path=csv_file_path,
            max_training=args.max_training,
            title=title,
            save_csv=not args.no_save_csv,
            save_plot=not args.no_save_plot
        )
        
        if result:
            print(f"成功处理 {len(result['valid_groups'])} 组有效数据")
    
    elif args.mode == 'compare':
        # 两类平均对比曲线模式
        if not os.path.exists(args.ir_file) and not os.path.exists(args.mr_file):
            print(f"错误: IR文件 '{args.ir_file}' 和 MR文件 '{args.mr_file}' 都不存在")
            exit(1)
        
        # 设置标题
        title = args.title if args.title else "IR组与MR组平均学习曲线对比"
        
        # 调用对比绘图函数
        result = plot_average_learning_curves(
            ir_csv_path=args.ir_file,
            mr_csv_path=args.mr_file,
            max_training=args.max_training,
            title=title,
            save_csv=not args.no_save_csv,
            save_plot=not args.no_save_plot
        )
        
        if result:
            print("成功生成IR组与MR组平均学习曲线对比图")
            for group_name, group_data in result.items():
                print(f"{group_name.upper()}组: {group_data['valid_count']} 个有效样本")
    
    print("程序执行完成")